import pymysql
from dotenv import load_dotenv
import os
import time
from contextlib import contextmanager

# 加载环境变量
load_dotenv()

def get_db_config():
    return {
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 3306)),
        'user': os.getenv('DB_USER', 'root'),
        'password': os.getenv('DB_PASSWORD', '123456'),
        'db': os.getenv('DB_NAME', 'fastapi'),
        'charset': 'utf8mb4'
    }

@contextmanager
def get_connection(max_retries=3, retry_delay=1):
    """获取数据库连接，带重试机制"""
    retries = 0
    last_error = None
    
    while retries < max_retries:
        try:
            connection = pymysql.connect(**get_db_config())
            print("MySQL 数据库连接成功")
            yield connection
            connection.close()
            print("数据库连接已关闭")
            return
        except Exception as e:
            last_error = e
            retries += 1
            if retries < max_retries:
                time.sleep(retry_delay)
    
    raise last_error 